package com.example.persistence;

import com.alibaba.csp.sentinel.command.handler.ModifyParamFlowRulesCommandHandler;
import com.alibaba.csp.sentinel.datasource.FileRefreshableDataSource;
import com.alibaba.csp.sentinel.datasource.FileWritableDataSource;
import com.alibaba.csp.sentinel.datasource.ReadableDataSource;
import com.alibaba.csp.sentinel.datasource.WritableDataSource;
import com.alibaba.csp.sentinel.init.InitFunc;
import com.alibaba.csp.sentinel.slots.block.authority.AuthorityRule;
import com.alibaba.csp.sentinel.slots.block.authority.AuthorityRuleManager;
import com.alibaba.csp.sentinel.slots.block.degrade.DegradeRule;
import com.alibaba.csp.sentinel.slots.block.degrade.DegradeRuleManager;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowRuleManager;
import com.alibaba.csp.sentinel.slots.system.SystemRule;
import com.alibaba.csp.sentinel.slots.system.SystemRuleManager;
import com.alibaba.csp.sentinel.transport.util.WritableDataSourceRegistry;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.List;
import java.util.Map;

public class PullModeByFileDataDataSource implements InitFunc {
    @Override
    public void init() throws Exception {

        mkdirIfNotExists(PersistenceConstant.ruleDir);

        createFileIfNotExists(PersistenceConstant.ruleMap);

        dealFlowRules();

        dealDegradeRules();

        dealSystemRules();

        dealParamFlowRules();

        dealAuthorityRules();

    }

    private void mkdirIfNotExists(String filePath) {
        File file = new File(filePath);
        if (!file.exists()) {
            file.mkdirs();
        }
    }

    private void createFileIfNotExists(Map<String, String> pathMap) throws IOException {
        for (Map.Entry<String, String> entry : pathMap.entrySet()) {
            File file = new File(entry.getValue());
            if (!file.exists()) {
                file.createNewFile();
            }
        }
    }

    private void dealFlowRules() throws FileNotFoundException {
        //创建流控规则的可读数据源
        ReadableDataSource<String, List<FlowRule>> ruleRDS = new FileRefreshableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.FLOW_RULE_PATH),
                s -> JSON.parseObject(s, new TypeReference<List<FlowRule>>() {
                })
        );

        // 将可读数据源注册至FlowRuleManager 这样当规则文件发生变化时，就会更新规则到内存
        FlowRuleManager.register2Property(ruleRDS.getProperty());

        WritableDataSource<List<FlowRule>> ruleWDS = new FileWritableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.FLOW_RULE_PATH),
                JSON::toJSONString
        );
        // 将可写数据源注册至 transport 模块的 WritableDataSourceRegistry 中.
        // 这样收到控制台推送的规则时，Sentinel 会先更新到内存，然后将规则写入到文件中.
        WritableDataSourceRegistry.registerFlowDataSource(ruleWDS);
    }

    private void dealDegradeRules() throws FileNotFoundException {
        ReadableDataSource<String, List<DegradeRule>> ruleRDS = new FileRefreshableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.DEGRADE_RULE_PATH),
                s -> JSON.parseObject(s, new TypeReference<List<DegradeRule>>() {
                })
        );

        DegradeRuleManager.register2Property(ruleRDS.getProperty());

        WritableDataSource<List<DegradeRule>> ruleWDS = new FileWritableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.DEGRADE_RULE_PATH),
                JSON::toJSONString
        );
        WritableDataSourceRegistry.registerDegradeDataSource(ruleWDS);
    }

    private void dealSystemRules() throws FileNotFoundException {
        ReadableDataSource<String, List<SystemRule>> ruleRDS = new FileRefreshableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.SYSTEM_RULE_PATH),
                s -> JSON.parseObject(s, new TypeReference<List<SystemRule>>() {
                })
        );

        SystemRuleManager.register2Property(ruleRDS.getProperty());

        WritableDataSource<List<SystemRule>> ruleWDS = new FileWritableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.SYSTEM_RULE_PATH),
                JSON::toJSONString
        );
        WritableDataSourceRegistry.registerSystemDataSource(ruleWDS);
    }

    private void dealAuthorityRules() throws FileNotFoundException {
        ReadableDataSource<String, List<AuthorityRule>> ruleRDS = new FileRefreshableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.AUTHORITY_RULE_PATH),
                s -> JSON.parseObject(s, new TypeReference<List<AuthorityRule>>() {
                })
        );

        AuthorityRuleManager.register2Property(ruleRDS.getProperty());

        WritableDataSource<List<AuthorityRule>> ruleWDS = new FileWritableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.AUTHORITY_RULE_PATH),
                JSON::toJSONString
        );
        WritableDataSourceRegistry.registerAuthorityDataSource(ruleWDS);
    }

    private void dealParamFlowRules() throws FileNotFoundException {
        ReadableDataSource<String, List<ParamFlowRule>> ruleRDS = new FileRefreshableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.PARAM_FLOW_RULE_PATH),
                s -> JSON.parseObject(s, new TypeReference<List<ParamFlowRule>>() {
                })
        );

        ParamFlowRuleManager.register2Property(ruleRDS.getProperty());

        WritableDataSource<List<ParamFlowRule>> ruleWDS = new FileWritableDataSource<>(
                PersistenceConstant.ruleMap.get(PersistenceConstant.PARAM_FLOW_RULE_PATH),
                JSON::toJSONString
        );
        ModifyParamFlowRulesCommandHandler.setWritableDataSource(ruleWDS);
    }
}
